tensor([[0.8400, 0.0000, 0.8400, 0.0200, 0.8400],
        [0.1600, 1.0000, 0.1600, 0.9800, 0.1600]], grad_fn=<RoundBackward1>)